/*
 * minimum_snap.cpp
 *
 *  Created on: 2022年1月27日
 *      Author: Huanyu Zhang
 */

#include "minimum_snap.hpp"

MinimumSnap minimumSnapObj(4, 3.0, 1.0); //order = 4，即p = 0, v = 1, a = 2, j = 3, s = 4

MinimumSnap::MinimumSnap(unsigned int order, double max_vel, double max_accel)
        : order_(order), max_vel_(max_vel), max_accel_(max_accel), num_poly_coeff(2 * order) {}


/*
    时间是通过计算梯形速度曲线而得到的
    设最大加速度为am，最大速度为vm，则v从0-vm的加速时间为t = vm / am
    则飞机从v = 0 - > vm - > 0所行驶的距离是：dist = 2 * 1/ 2 * am * t * t = am * t * t
    1. 当两点之间的距离d1>d时
        飞机剩下的距离应该是以最大速度进行行驶，即t1 = 2 * t + (d1 - d) / vm
    2. 当两点之间的距离d1 <= d时
        应该将d1分开两段，一段加速，一段减速，即1 / 2 * d1 = 1 / 2 * am * t0 * t0
        即t0 = sqrt(d1 / am)，但是注意t0只是加速或者减速的时间，因此两段应该是2 * t0
*/
VecXd MinimumSnap::AllocateTime(const MatXd &waypoint) const {
    VecXd times = VecXd::Zero(waypoint.rows() - 1);
    const double t = max_vel_ / max_accel_;
    const double dist_threshold_1 = max_accel_ * t * t;

    double segment_t;
    for (int i = 1; i < waypoint.rows(); ++i) {
        double delta_dist = (waypoint.row(i) - waypoint.row(i - 1)).norm();
        if (delta_dist > dist_threshold_1) {
            segment_t = t * 2 + (delta_dist - dist_threshold_1) / max_vel_;
        } else {
            segment_t = 2 * std::sqrt(delta_dist / max_accel_);
        }
        times[i - 1] = segment_t;
    }
    return times;
}

unsigned int MinimumSnap::GetPolyCoeffNum() const {
    return 2u * order_;
}

MatrixXd MinimumSnap::polydiff(MatrixXd& poly_coeff_mat) {
    const unsigned int num_segments = poly_coeff_mat.rows();
    int cur_poly_num = poly_coeff_mat.cols() - 1; //每求一次导数，其系数个数就会减少一个
    MatrixXd poly_coeff_diff_mat = MatrixXd::Zero(num_segments, cur_poly_num);
    for (unsigned int i = 0; i < num_segments; ++i) {
        for (unsigned int j = 0; j < cur_poly_num; ++j) {
            //if (j = num_poly_coeff - 1) continue; //0阶多项式的微分为0
            poly_coeff_diff_mat(i, j) = poly_coeff_mat(i, j) * (cur_poly_num - j);
        }
    }
    return poly_coeff_diff_mat;
}

/*
    yaw∈[-pi, pi]
    yaw = dy / dx = (dy / dt) / (dx / dt) = vy(t) / vx(t)
	因此atan2(vy(t), vx(t)) -->转换成[-pi, pi]就是yaw了
*/
void MinimumSnap::getCommandPosVelAccYaw(double* res, const MatrixXd& poly_coeff_mat_x, const MatrixXd& poly_coeff_mat_y, const MatrixXd& poly_coeff_diff_mat_x, const MatrixXd& poly_coeff_diff_mat_y, const MatrixXd& poly_coeff_acc_mat_x, const MatrixXd& poly_coeff_acc_mat_y, int k, double t) {
    VectorXd coeff_x = poly_coeff_mat_x.row(k);
    VectorXd coeff_y = poly_coeff_mat_y.row(k);
    VectorXd coeff_vx = poly_coeff_diff_mat_x.row(k);
    VectorXd coeff_vy = poly_coeff_diff_mat_y.row(k);
    VectorXd coeff_ax = poly_coeff_acc_mat_x.row(k);
    VectorXd coeff_ay = poly_coeff_acc_mat_y.row(k);
    VectorXd time = VectorXd::Zero(num_poly_coeff);

    for (unsigned int i = 0; i < num_poly_coeff; ++i) {
        if (i == 0) {
            time(i) = 1.0;
        }
        else {
            time(i) = pow(t, i);
        }
    }
    double temp_x = 0.0;
    double temp_y = 0.0;
    double temp_vel_x = 0.0;
    double temp_vel_y = 0.0;
    double temp_acc_x = 0.0;
    double temp_acc_y = 0.0;
    for (int i = 0u; i < time.rows(); ++i) {
    	temp_x = temp_x + coeff_x(i) * time(time.rows() - i - 1u);
    	temp_y = temp_y + coeff_y(i) * time(time.rows() - i - 1u);
    	if(i < time.rows() - 1){ //速度项只有7个参数，因此最后一个必须舍弃
            temp_vel_x = temp_vel_x + coeff_vx(i) * time(time.rows() - i - 2u);
            temp_vel_y = temp_vel_y + coeff_vy(i) * time(time.rows() - i - 2u);
    	}
    	if(i < time.rows() - 2){ //加速度项只有6个参数，因此最后两个必须舍弃
    		temp_acc_x = temp_acc_x + coeff_ax(i) * time(time.rows() - i - 3u);
            temp_acc_y = temp_acc_y + coeff_ay(i) * time(time.rows() - i - 3u);
    	}
    }
    double tempYaw = atan2(temp_vel_y, temp_vel_x);
    res[0] = temp_x;
    res[1] = temp_y;
    res[2] = temp_vel_x;
    res[3] = temp_vel_y;
    res[4] = temp_acc_x;
	res[5] = temp_acc_y;
    res[6] = tempYaw;
    return;
}

MatXd MinimumSnap::SolveQPClosedForm(const VecXd &waypoints_pos, const VecXd &waypoints_vel,
                                     const VecXd &waypoints_acc, const VecXd &segments_time) const {

    const unsigned int poly_order = 2 * order_ - 1; //7阶
    const unsigned int num_poly_coeff = poly_order + 1; //8个参数
    const unsigned int num_segments = segments_time.size();

    const unsigned int num_all_poly_coeff = num_poly_coeff * num_segments;

    const unsigned int A_block_rows = 2u * order_; //8
    const unsigned int A_block_cols = num_poly_coeff; //8
    MatXd A = MatXd::Zero(num_segments * A_block_rows, num_segments * A_block_cols); //微分约束矩阵

    for (unsigned int i = 0; i < num_segments; ++i) {
        unsigned int row = i * A_block_rows;
        unsigned int col = i * A_block_cols;

        MatXd sub_A = MatXd::Zero(A_block_rows, A_block_cols);

        for (unsigned int j = 0; j < order_; ++j) { //代表p v a j
            for (unsigned int k = 0; k < num_poly_coeff; ++k) { //代表某个项式的阶数
                if (k < j) { //位置p对应的8个系数都不为0，v对应的p0系数为0...
                    continue;
                }

                sub_A(j, num_poly_coeff - 1 - k) = Factorial(k) / Factorial(k - j) * std::pow(0, k - j); //A矩阵次数高的放在前面 ,第i段轨迹的起点， t = 0
                sub_A(j + order_, num_poly_coeff - 1 - k) =
                        Factorial(k) / Factorial(k - j) * pow(segments_time[i], k - j); //第i段轨迹的终点，t = segments_time[i]
            }
        }

        A.block(row, col, A_block_rows, A_block_cols) = sub_A;
    }

    const unsigned int num_valid_variables = (num_segments + 1) * order_; //这里将连续性约束也加上去了，因为原来一共有2 * num_segements * order_
    const unsigned int num_fixed_variables = 2u * order_ + (num_segments - 1); //固定变量，相当于df

    MatXd C_T = MatXd::Zero(num_all_poly_coeff, num_valid_variables);
    for (unsigned int i = 0; i < num_all_poly_coeff; ++i) {
        if (i < order_) { //前面4个
            C_T(i, i) = 1.0;
            continue;
        }

        if (i >= num_all_poly_coeff - order_) { //后面4个
            const unsigned int delta_index = i - (num_all_poly_coeff - order_);
            C_T(i, num_fixed_variables - order_ + delta_index) = 1.0;
            continue;
        }

        if ((i % order_ == 0u) && (i / order_ % 2u == 1u)) {
            const unsigned int index = i / (2u * order_) + order_;
            C_T(i, index) = 1.0;
            continue;
        }

        if ((i % order_ == 0u) && (i / order_ % 2u == 0u)) {
            const unsigned int index = i / (2u * order_) + order_ - 1u;
            C_T(i, index) = 1.0;
            continue;
        }

        if ((i % order_ != 0u) && (i / order_ % 2u == 1u)) {
            const unsigned int temp_index_0 = i / (2 * order_) * (2 * order_) + order_;
            const unsigned int temp_index_1 = i / (2 * order_) * (order_ - 1) + i - temp_index_0 - 1;
            C_T(i, num_fixed_variables + temp_index_1) = 1.0;
            continue;
        }

        if ((i % order_ != 0u) && (i / order_ % 2u == 0u)) {
            const unsigned int temp_index_0 = (i - order_) / (2 * order_) * (2 * order_) + order_;
            const unsigned int temp_index_1 =
                    (i - order_) / (2 * order_) * (order_ - 1) + (i - order_) - temp_index_0 - 1;
            C_T(i, num_fixed_variables + temp_index_1) = 1.0;
            continue;
        }
    }

    MatXd Q = MatXd::Zero(num_all_poly_coeff, num_all_poly_coeff); // (8 * seg) * (8 * seg)
    for (unsigned int k = 0u; k < num_segments; ++k) {
        MatXd sub_Q = MatXd::Zero(num_poly_coeff, num_poly_coeff); //8 * 8
        for (unsigned int i = 0u; i <= poly_order; ++i) {
            for (unsigned int l = 0u; l <= poly_order; ++l) {
                if (num_poly_coeff - i <= order_ || num_poly_coeff - l <= order_) {
                    continue;
                }

                sub_Q(i, l) = (Factorial(poly_order - i) / Factorial(poly_order - order_ - i)) *
                              (Factorial(poly_order - l) / Factorial(poly_order - order_ - l)) /
                              (poly_order - i + poly_order - l - (2 * order_ - 1)) *
                              std::pow(segments_time[k], poly_order - i + poly_order - l - (2 * order_ - 1));
            }
        }

        const unsigned int row = k * num_poly_coeff;
        Q.block(row, row, num_poly_coeff, num_poly_coeff) = sub_Q;
    }

    MatXd R = C_T.transpose() * A.transpose().inverse() * Q * A.inverse() * C_T;

    VecXd d_selected = VecXd::Zero(num_valid_variables); // d_selected = [df; dp];
    for (unsigned int i = 0; i < num_all_poly_coeff; ++i) {
        if (i == 0u) { //df的前3个 p
            d_selected[i] = waypoints_pos[0];
            continue;
        }

        if (i == 1u && order_ >= 2u) { //v
            d_selected[i] = waypoints_vel[0];
            continue;
        }

        if (i == 2u && order_ >= 3u) { //a
            d_selected[i] = waypoints_acc[0];
            continue;
        }

        if (i == num_all_poly_coeff - order_ + 2u && order_ >= 3u) { //df的最后3个pva，这里如果严格按照顺序，应该是waypoints_acc[num_segments]
            d_selected(num_fixed_variables - order_ + 2) = waypoints_acc[1];
            continue;
        }

        if (i == num_all_poly_coeff - order_ + 1u && order_ >= 2u) {
            d_selected(num_fixed_variables - order_ + 1) = waypoints_vel[1];
            continue;
        }

        if (i == num_all_poly_coeff - order_) {
            d_selected(num_fixed_variables - order_) = waypoints_pos[num_segments];
            continue;
        }

        if ((i % order_ == 0u) && (i / order_ % 2u == 0u)) { //中间的waypoint
            const unsigned int index = i / (2 * order_) + order_ - 1; //由于前面已经占用了pvaj，所以要添加order_偏移量，由于第一个满足2 * order_下标应该是0，所以应该-1
            d_selected(index) = waypoints_pos[i / (2 * order_)];
            continue;
        }
    }

    MatXd R_PP = R.block(num_fixed_variables, num_fixed_variables,
                         num_valid_variables - num_fixed_variables,
                         num_valid_variables - num_fixed_variables);

    VecXd d_F = d_selected.head(num_fixed_variables);
    MatXd R_FP = R.block(0, num_fixed_variables, num_fixed_variables,
                         num_valid_variables - num_fixed_variables);

    MatXd d_optimal = -R_PP.inverse() * R_FP.transpose() * d_F;

    d_selected.tail(num_valid_variables - num_fixed_variables) = d_optimal;
    VecXd d = C_T * d_selected;

    VecXd P = A.inverse() * d;

    MatXd poly_coeff_mat = MatXd::Zero(num_segments, num_poly_coeff);
    for (unsigned int i = 0; i < num_segments; ++i) {
        poly_coeff_mat.block(i, 0, 1, num_poly_coeff) =
                P.block(num_poly_coeff * i, 0, num_poly_coeff, 1).transpose();
    }

    return poly_coeff_mat;
}
